"""
AI Historical Simulation Platform - API Server

This module provides a comprehensive RESTful API server with WebSocket support,
authentication, rate limiting, and resource management for the AI Historical
Simulation Platform.
"""

import asyncio
import json
import logging
import time
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Set
from dataclasses import dataclass, asdict
import hashlib
import secrets
from contextlib import asynccontextmanager

try:
    from fastapi import FastAPI, HTTPException, Depends, WebSocket, WebSocketDisconnect
    from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
    from fastapi.middleware.cors import CORSMiddleware
    from fastapi.middleware.trustedhost import TrustedHostMiddleware
    from fastapi.responses import JSONResponse
    from pydantic import BaseModel, Field
    import uvicorn
    FASTAPI_AVAILABLE = True
except ImportError:
    FASTAPI_AVAILABLE = False

from .session_manager import SessionManager, get_or_create_session

logger = logging.getLogger(__name__)


# Pydantic models for API endpoints
class ChatMessage(BaseModel):
    """Chat message model."""
    figure_name: str = Field(..., description="Name of historical figure")
    message: str = Field(..., max_length=2000, description="User message")
    session_id: Optional[str] = Field(None, description="Optional session ID")
    context: Optional[Dict[str, Any]] = Field(default_factory=dict, description="Additional context")


class ChatResponse(BaseModel):
    """Chat response model."""
    figure_name: str
    response: str
    session_id: str
    response_time_ms: float
    metadata: Dict[str, Any] = Field(default_factory=dict)


class SessionCreate(BaseModel):
    """Session creation model."""
    user_id: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)
    priority_level: int = Field(1, ge=1, le=3, description="Session priority (1-3)")


class CrossTemporalConversation(BaseModel):
    """Cross-temporal conversation model."""
    figure_names: List[str] = Field(..., min_items=2, max_items=5)
    topic: str = Field(..., max_length=500)
    rounds: int = Field(3, ge=1, le=10)
    session_id: Optional[str] = None


class UserAuth(BaseModel):
    """User authentication model."""
    username: str
    password: str


class ApiKey(BaseModel):
    """API key model."""
    key: str
    user_id: str
    permissions: List[str]
    expires_at: Optional[datetime] = None
    rate_limit: int = 100  # requests per minute


@dataclass
class RateLimitInfo:
    """Rate limit tracking information."""
    requests: int
    window_start: datetime
    max_requests: int
    window_duration: timedelta


class AuthenticationManager:
    """Manages API authentication and authorization."""
    
    def __init__(self):
        self.api_keys: Dict[str, ApiKey] = {}
        self.user_sessions: Dict[str, str] = {}  # user_id -> session_token
        self.rate_limits: Dict[str, RateLimitInfo] = {}
        
        # Create default admin key for development
        admin_key = self._generate_api_key()
        self.api_keys[admin_key] = ApiKey(
            key=admin_key,
            user_id="admin",
            permissions=["read", "write", "admin"],
            rate_limit=1000
        )
        logger.info(f"Generated admin API key: {admin_key}")
    
    def _generate_api_key(self) -> str:
        """Generate secure API key."""
        return f"hsp_{secrets.token_urlsafe(32)}"
    
    def create_api_key(self, 
                      user_id: str, 
                      permissions: List[str], 
                      rate_limit: int = 100,
                      expires_at: Optional[datetime] = None) -> str:
        """Create new API key."""
        key = self._generate_api_key()
        self.api_keys[key] = ApiKey(
            key=key,
            user_id=user_id,
            permissions=permissions,
            expires_at=expires_at,
            rate_limit=rate_limit
        )
        return key
    
    def validate_api_key(self, key: str) -> Optional[ApiKey]:
        """Validate API key and return user info."""
        api_key_obj = self.api_keys.get(key)
        
        if not api_key_obj:
            return None
        
        # Check expiration
        if api_key_obj.expires_at and datetime.now() > api_key_obj.expires_at:
            return None
        
        return api_key_obj
    
    def check_rate_limit(self, api_key: str, user_id: str) -> bool:
        """Check if user is within rate limits."""
        key_info = self.api_keys.get(api_key)
        if not key_info:
            return False
        
        now = datetime.now()
        rate_limit_key = f"{api_key}:{user_id}"
        
        # Get or create rate limit info
        if rate_limit_key not in self.rate_limits:
            self.rate_limits[rate_limit_key] = RateLimitInfo(
                requests=0,
                window_start=now,
                max_requests=key_info.rate_limit,
                window_duration=timedelta(minutes=1)
            )
        
        rate_info = self.rate_limits[rate_limit_key]
        
        # Reset window if expired
        if now - rate_info.window_start > rate_info.window_duration:
            rate_info.requests = 0
            rate_info.window_start = now
        
        # Check limit
        if rate_info.requests >= rate_info.max_requests:
            return False
        
        # Increment counter
        rate_info.requests += 1
        return True
    
    def has_permission(self, api_key: str, permission: str) -> bool:
        """Check if API key has specific permission."""
        key_info = self.api_keys.get(api_key)
        if not key_info:
            return False
        
        return permission in key_info.permissions or "admin" in key_info.permissions


class WebSocketManager:
    """Manages WebSocket connections and real-time communication."""
    
    def __init__(self, platform):
        self.platform = platform
        self.active_connections: Dict[str, WebSocket] = {}
        self.connection_sessions: Dict[str, str] = {}  # connection_id -> session_id
        self.session_connections: Dict[str, Set[str]] = {}  # session_id -> connection_ids
    
    async def connect(self, websocket: WebSocket, connection_id: str) -> None:
        """Accept WebSocket connection."""
        await websocket.accept()
        self.active_connections[connection_id] = websocket
        logger.info(f"WebSocket connected: {connection_id}")
    
    def disconnect(self, connection_id: str) -> None:
        """Remove WebSocket connection."""
        if connection_id in self.active_connections:
            del self.active_connections[connection_id]
        
        # Clean up session associations
        if connection_id in self.connection_sessions:
            session_id = self.connection_sessions[connection_id]
            del self.connection_sessions[connection_id]
            
            if session_id in self.session_connections:
                self.session_connections[session_id].discard(connection_id)
                if not self.session_connections[session_id]:
                    del self.session_connections[session_id]
        
        logger.info(f"WebSocket disconnected: {connection_id}")
    
    def associate_session(self, connection_id: str, session_id: str) -> None:
        """Associate WebSocket connection with session."""
        self.connection_sessions[connection_id] = session_id
        
        if session_id not in self.session_connections:
            self.session_connections[session_id] = set()
        self.session_connections[session_id].add(connection_id)
    
    async def send_to_connection(self, connection_id: str, message: Dict[str, Any]) -> bool:
        """Send message to specific connection."""
        if connection_id in self.active_connections:
            try:
                await self.active_connections[connection_id].send_text(json.dumps(message))
                return True
            except Exception as e:
                logger.error(f"Failed to send to {connection_id}: {e}")
                self.disconnect(connection_id)
        return False
    
    async def send_to_session(self, session_id: str, message: Dict[str, Any]) -> int:
        """Send message to all connections in a session."""
        if session_id not in self.session_connections:
            return 0
        
        sent_count = 0
        connections = list(self.session_connections[session_id])
        
        for connection_id in connections:
            if await self.send_to_connection(connection_id, message):
                sent_count += 1
        
        return sent_count
    
    async def broadcast(self, message: Dict[str, Any]) -> int:
        """Broadcast message to all connections."""
        sent_count = 0
        connections = list(self.active_connections.keys())
        
        for connection_id in connections:
            if await self.send_to_connection(connection_id, message):
                sent_count += 1
        
        return sent_count


class APIServer:
    """
    Production-ready API server for the AI Historical Simulation Platform.
    
    Features:
    - RESTful API endpoints for all platform functionality
    - WebSocket support for real-time communication
    - Authentication and authorization
    - Rate limiting and resource management
    - Comprehensive error handling
    - API documentation and monitoring
    """
    
    def __init__(self, platform, host: str = "0.0.0.0", port: int = 8000):
        """
        Initialize API server.
        
        Args:
            platform: AI Historical Simulation Platform instance
            host: Server host address
            port: Server port
        """
        if not FASTAPI_AVAILABLE:
            raise ImportError("FastAPI not available. Install with: pip install fastapi uvicorn")
        
        self.platform = platform
        self.host = host
        self.port = port
        
        # Initialize components
        self.auth_manager = AuthenticationManager()
        self.ws_manager = WebSocketManager(platform)
        
        # Create FastAPI app
        self.app = self._create_app()
        
        # Server state
        self.is_running = False
        self.server_task: Optional[asyncio.Task] = None
        
        logger.info(f"APIServer initialized on {host}:{port}")
    
    def _create_app(self) -> FastAPI:
        """Create and configure FastAPI application."""
        app = FastAPI(
            title="AI Historical Simulation Platform API",
            description="Comprehensive API for interacting with historical figures through advanced AI simulation",
            version="1.0.0",
            docs_url="/docs",
            redoc_url="/redoc"
        )
        
        # Add middleware
        app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],  # Configure appropriately for production
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )
        
        app.add_middleware(
            TrustedHostMiddleware,
            allowed_hosts=["*"]  # Configure appropriately for production
        )
        
        # Store references
        app.state.platform = self.platform
        app.state.auth_manager = self.auth_manager
        app.state.ws_manager = self.ws_manager
        
        # Setup routes
        self._setup_routes(app)
        
        return app
    
    def _setup_routes(self, app: FastAPI) -> None:
        """Setup all API routes."""
        
        # Authentication dependency
        security = HTTPBearer()
        
        async def authenticate(credentials: HTTPAuthorizationCredentials = Depends(security)) -> ApiKey:
            """Authentication dependency."""
            api_key_obj = self.auth_manager.validate_api_key(credentials.credentials)
            if not api_key_obj:
                raise HTTPException(status_code=401, detail="Invalid API key")
            
            # Check rate limit
            if not self.auth_manager.check_rate_limit(credentials.credentials, api_key_obj.user_id):
                raise HTTPException(status_code=429, detail="Rate limit exceeded")
            
            return api_key_obj
        
        # Health check endpoint
        @app.get("/health")
        async def health_check():
            """System health check."""
            try:
                health = await self.platform.get_platform_health()
                return {
                    "status": "healthy" if health.is_healthy else "unhealthy",
                    "health": asdict(health),
                    "timestamp": datetime.now().isoformat()
                }
            except Exception as e:
                return JSONResponse(
                    status_code=503,
                    content={
                        "status": "error",
                        "error": str(e),
                        "timestamp": datetime.now().isoformat()
                    }
                )
        
        # Platform metrics endpoint
        @app.get("/metrics")
        async def get_metrics(auth: ApiKey = Depends(authenticate)):
            """Get comprehensive platform metrics."""
            if not self.auth_manager.has_permission(auth.key, "read"):
                raise HTTPException(status_code=403, detail="Insufficient permissions")
            
            metrics = await self.platform.get_system_metrics()
            return {"metrics": metrics, "timestamp": datetime.now().isoformat()}
        
        # Session management endpoints
        @app.post("/sessions", response_model=dict)
        async def create_session(session_data: SessionCreate, auth: ApiKey = Depends(authenticate)):
            """Create new user session."""
            if not self.auth_manager.has_permission(auth.key, "write"):
                raise HTTPException(status_code=403, detail="Insufficient permissions")
            
            try:
                session_id = await self.platform.session_manager.create_session(
                    user_id=session_data.user_id or auth.user_id,
                    session_metadata=session_data.metadata,
                    priority_level=session_data.priority_level
                )
                
                return {
                    "session_id": session_id,
                    "created_at": datetime.now().isoformat(),
                    "status": "created"
                }
            except Exception as e:
                raise HTTPException(status_code=400, detail=str(e))
        
        @app.get("/sessions/{session_id}")
        async def get_session(session_id: str, auth: ApiKey = Depends(authenticate)):
            """Get session information."""
            session = await self.platform.session_manager.get_session(session_id)
            if not session:
                raise HTTPException(status_code=404, detail="Session not found")
            
            # Check permission to access session
            if session.user_id != auth.user_id and not self.auth_manager.has_permission(auth.key, "admin"):
                raise HTTPException(status_code=403, detail="Access denied")
            
            return {"session": session.to_dict()}
        
        @app.delete("/sessions/{session_id}")
        async def delete_session(session_id: str, auth: ApiKey = Depends(authenticate)):
            """Delete session."""
            session = await self.platform.session_manager.get_session(session_id)
            if not session:
                raise HTTPException(status_code=404, detail="Session not found")
            
            # Check permission to delete session
            if session.user_id != auth.user_id and not self.auth_manager.has_permission(auth.key, "admin"):
                raise HTTPException(status_code=403, detail="Access denied")
            
            success = await self.platform.session_manager.remove_session(session_id)
            return {"deleted": success}
        
        # Historical figure endpoints
        @app.get("/figures")
        async def list_figures(auth: ApiKey = Depends(authenticate)):
            """List available historical figures."""
            figure_names = list(self.platform.historical_figures.keys())
            return {"figures": figure_names, "count": len(figure_names)}
        
        @app.get("/figures/{figure_name}")
        async def get_figure_info(figure_name: str, auth: ApiKey = Depends(authenticate)):
            """Get information about a historical figure."""
            if figure_name not in self.platform.historical_figures:
                # Try to load the figure
                try:
                    await self.platform.load_historical_figure(figure_name)
                except Exception:
                    raise HTTPException(status_code=404, detail="Historical figure not found")
            
            figure = self.platform.historical_figures[figure_name]
            return {"figure": figure.to_dict()}
        
        @app.post("/figures/{figure_name}/load")
        async def load_figure(figure_name: str, auth: ApiKey = Depends(authenticate)):
            """Load a historical figure."""
            if not self.auth_manager.has_permission(auth.key, "write"):
                raise HTTPException(status_code=403, detail="Insufficient permissions")
            
            try:
                figure = await self.platform.load_historical_figure(figure_name, force_reload=True)
                return {
                    "figure": figure.to_dict(),
                    "status": "loaded",
                    "timestamp": datetime.now().isoformat()
                }
            except Exception as e:
                raise HTTPException(status_code=400, detail=str(e))
        
        # Chat endpoints
        @app.post("/chat", response_model=ChatResponse)
        async def chat_with_figure(message: ChatMessage, auth: ApiKey = Depends(authenticate)):
            """Chat with a historical figure."""
            if not self.auth_manager.has_permission(auth.key, "write"):
                raise HTTPException(status_code=403, detail="Insufficient permissions")
            
            try:
                start_time = time.time()
                
                # Get or create session
                session_id, session = await get_or_create_session(
                    self.platform.session_manager,
                    session_id=message.session_id,
                    user_id=auth.user_id
                )
                
                # Generate response
                response = await self.platform.generate_response(
                    session_id=session_id,
                    figure_name=message.figure_name,
                    user_input=message.message,
                    context=message.context
                )
                
                response_time = (time.time() - start_time) * 1000
                
                return ChatResponse(
                    figure_name=message.figure_name,
                    response=response,
                    session_id=session_id,
                    response_time_ms=response_time,
                    metadata={
                        "timestamp": datetime.now().isoformat(),
                        "user_id": auth.user_id
                    }
                )
                
            except Exception as e:
                raise HTTPException(status_code=400, detail=str(e))
        
        # Cross-temporal conversation endpoint
        @app.post("/conversations/cross-temporal")
        async def create_cross_temporal_conversation(
            conversation: CrossTemporalConversation, 
            auth: ApiKey = Depends(authenticate)
        ):
            """Create cross-temporal conversation between historical figures."""
            if not self.auth_manager.has_permission(auth.key, "write"):
                raise HTTPException(status_code=403, detail="Insufficient permissions")
            
            try:
                # Get or create session
                session_id, _ = await get_or_create_session(
                    self.platform.session_manager,
                    session_id=conversation.session_id,
                    user_id=auth.user_id
                )
                
                # Create conversation
                conversation_result = await self.platform.create_cross_temporal_conversation(
                    session_id=session_id,
                    figure_names=conversation.figure_names,
                    topic=conversation.topic,
                    rounds=conversation.rounds
                )
                
                return {
                    "conversation": conversation_result,
                    "session_id": session_id,
                    "timestamp": datetime.now().isoformat()
                }
                
            except Exception as e:
                raise HTTPException(status_code=400, detail=str(e))
        
        # WebSocket endpoint
        @app.websocket("/ws/{connection_id}")
        async def websocket_endpoint(websocket: WebSocket, connection_id: str):
            """WebSocket endpoint for real-time communication."""
            await self.ws_manager.connect(websocket, connection_id)
            
            try:
                while True:
                    # Receive message
                    data = await websocket.receive_text()
                    message = json.loads(data)
                    
                    await self._handle_websocket_message(connection_id, message)
                    
            except WebSocketDisconnect:
                self.ws_manager.disconnect(connection_id)
            except Exception as e:
                logger.error(f"WebSocket error for {connection_id}: {e}")
                self.ws_manager.disconnect(connection_id)
        
        # Admin endpoints
        @app.post("/admin/api-keys")
        async def create_api_key(
            key_data: dict,
            auth: ApiKey = Depends(authenticate)
        ):
            """Create new API key (admin only)."""
            if not self.auth_manager.has_permission(auth.key, "admin"):
                raise HTTPException(status_code=403, detail="Admin access required")
            
            new_key = self.auth_manager.create_api_key(
                user_id=key_data.get("user_id"),
                permissions=key_data.get("permissions", ["read"]),
                rate_limit=key_data.get("rate_limit", 100),
                expires_at=key_data.get("expires_at")
            )
            
            return {"api_key": new_key}
        
        @app.get("/admin/sessions")
        async def list_all_sessions(auth: ApiKey = Depends(authenticate)):
            """List all active sessions (admin only)."""
            if not self.auth_manager.has_permission(auth.key, "admin"):
                raise HTTPException(status_code=403, detail="Admin access required")
            
            session_ids = await self.platform.session_manager.get_active_sessions()
            sessions = []
            
            for session_id in session_ids:
                session = await self.platform.session_manager.get_session(session_id)
                if session:
                    sessions.append(session.to_dict())
            
            return {"sessions": sessions, "count": len(sessions)}
    
    async def _handle_websocket_message(self, connection_id: str, message: Dict[str, Any]) -> None:
        """Handle incoming WebSocket message."""
        try:
            msg_type = message.get("type")
            
            if msg_type == "auth":
                # Authenticate WebSocket connection
                api_key = message.get("api_key")
                if not api_key:
                    await self.ws_manager.send_to_connection(
                        connection_id, 
                        {"type": "error", "message": "API key required"}
                    )
                    return
                
                key_info = self.auth_manager.validate_api_key(api_key)
                if not key_info:
                    await self.ws_manager.send_to_connection(
                        connection_id,
                        {"type": "error", "message": "Invalid API key"}
                    )
                    return
                
                await self.ws_manager.send_to_connection(
                    connection_id,
                    {"type": "auth_success", "user_id": key_info.user_id}
                )
            
            elif msg_type == "join_session":
                # Associate connection with session
                session_id = message.get("session_id")
                if session_id:
                    self.ws_manager.associate_session(connection_id, session_id)
                    await self.ws_manager.send_to_connection(
                        connection_id,
                        {"type": "session_joined", "session_id": session_id}
                    )
            
            elif msg_type == "chat":
                # Handle chat message via WebSocket
                await self._handle_websocket_chat(connection_id, message)
            
            elif msg_type == "get_figures":
                # Send list of available figures
                figures = list(self.platform.historical_figures.keys())
                await self.ws_manager.send_to_connection(
                    connection_id,
                    {"type": "figures_list", "figures": figures}
                )
            
            else:
                await self.ws_manager.send_to_connection(
                    connection_id,
                    {"type": "error", "message": f"Unknown message type: {msg_type}"}
                )
                
        except Exception as e:
            logger.error(f"WebSocket message handling error: {e}")
            await self.ws_manager.send_to_connection(
                connection_id,
                {"type": "error", "message": "Message processing failed"}
            )
    
    async def _handle_websocket_chat(self, connection_id: str, message: Dict[str, Any]) -> None:
        """Handle chat message via WebSocket."""
        try:
            figure_name = message.get("figure")
            user_message = message.get("message")
            session_id = self.ws_manager.connection_sessions.get(connection_id)
            
            if not all([figure_name, user_message, session_id]):
                await self.ws_manager.send_to_connection(
                    connection_id,
                    {"type": "error", "message": "Missing required fields"}
                )
                return
            
            # Generate response
            start_time = time.time()
            response = await self.platform.generate_response(
                session_id=session_id,
                figure_name=figure_name,
                user_input=user_message
            )
            response_time = (time.time() - start_time) * 1000
            
            # Send response
            await self.ws_manager.send_to_connection(
                connection_id,
                {
                    "type": "chat_response",
                    "figure": figure_name,
                    "response": response,
                    "response_time_ms": response_time,
                    "timestamp": datetime.now().isoformat()
                }
            )
            
        except Exception as e:
            await self.ws_manager.send_to_connection(
                connection_id,
                {"type": "error", "message": f"Chat processing failed: {str(e)}"}
            )
    
    async def start(self) -> None:
        """Start the API server."""
        if self.is_running:
            return
        
        logger.info(f"Starting API server on {self.host}:{self.port}")
        
        config = uvicorn.Config(
            self.app,
            host=self.host,
            port=self.port,
            log_level="info",
            access_log=True
        )
        
        server = uvicorn.Server(config)
        self.server_task = asyncio.create_task(server.serve())
        self.is_running = True
        
        await self.server_task
    
    async def stop(self) -> None:
        """Stop the API server."""
        if not self.is_running:
            return
        
        logger.info("Stopping API server")
        
        if self.server_task:
            self.server_task.cancel()
            try:
                await self.server_task
            except asyncio.CancelledError:
                pass
        
        self.is_running = False
        logger.info("API server stopped")
    
    def get_admin_api_key(self) -> str:
        """Get admin API key for testing/setup."""
        for key, key_info in self.auth_manager.api_keys.items():
            if key_info.user_id == "admin":
                return key
        return ""


# Context manager for server lifecycle
@asynccontextmanager
async def managed_api_server(platform, **kwargs):
    """Context manager for API server lifecycle."""
    server = APIServer(platform, **kwargs)
    try:
        # Start server in background
        server_task = asyncio.create_task(server.start())
        yield server
    finally:
        await server.stop()


# Utility functions
def create_development_server(platform, host: str = "127.0.0.1", port: int = 8000) -> APIServer:
    """Create API server configured for development."""
    return APIServer(platform, host=host, port=port)


def create_production_server(platform, host: str = "0.0.0.0", port: int = 8000) -> APIServer:
    """Create API server configured for production."""
    server = APIServer(platform, host=host, port=port)
    
    # Add production-specific configurations here
    # - TLS configuration
    # - Security headers
    # - Rate limiting enhancements
    # - Monitoring integration
    
    return server